Make FirstValue an UDAF, Change `AggregateUDFImpl::accumulator` signature, support ORDER BY for UDAFs #9874

jayzhan211 · 2024-03-30T14:30:26Z

Which issue does this PR close?

First step of #8708
Close #9249

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Signed-off-by: jayzhan211 <[email protected]>

datafusion/core/src/execution/context/mod.rs

datafusion/core/src/physical_planner.rs

Signed-off-by: jayzhan211 <[email protected]>

alamb

❤️ -- looking very good @jayzhan211

alamb · 2024-04-01T20:47:44Z

datafusion/expr/src/expr_fn.rs

    }

    fn state_type(&self, _return_type: &DataType) -> Result<Vec<DataType>> {
        Ok(self.state_type.clone())
    }
 }

+pub struct FirstValue {


I see -- it seems like the issue is that the accumulator implementation requires PhysicalSortExpr.

To pull the code into its own crate maybe we could pull out the relevant pieces of datafusion-physical-expr into datafusion-physical-core or something (as a follow on PR)

alamb · 2024-04-01T20:49:17Z

datafusion/expr/src/expr_fn.rs

@@ -710,6 +712,16 @@ pub fn create_udaf(
    ))
 }

+/// Creates a new UDAF with a specific signature, state type and return type.
+/// The signature and state type must match the `Accumulator's implementation`.


It might make sense to add a comment explaining this is a temporary solution (i.e. that the idea is we'll pull the function out into its own crate, but for now we need to keep the physical implementation separate

alamb · 2024-04-01T20:49:40Z

datafusion/expr/src/expr_fn.rs

+        f.debug_struct("FirstValue")
+            .field("name", &self.name)
+            .field("signature", &self.signature)
+            .field("fun", &"<FUNC>")


Suggested change

.field("fun", &"<FUNC>")

.field("accumulator", &"<FUNC>")

alamb · 2024-04-01T20:50:48Z

datafusion/expr/src/function.rs

+    pub data_type: &'a DataType, // the return type of the function
+    pub schema: &'a Schema,      // the schema of the input arguments
+    pub ignore_nulls: bool,      // whether to ignore nulls


Since these field are pub I think they show up in the rustdocs. Can you please make them with three /// and above the field in question to show up in the docs?

alamb · 2024-04-01T20:52:04Z

datafusion/expr/src/function.rs

+    pub schema: &'a Schema,      // the schema of the input arguments
+    pub ignore_nulls: bool,      // whether to ignore nulls
+
+    // ordering arguments


Can you also please document how we would tell if there was no ORDER BY specified? (is sort_exprs empty?)

alamb · 2024-04-01T20:52:21Z

datafusion/expr/src/function.rs

+    pub ignore_nulls: bool,      // whether to ignore nulls
+
+    // ordering arguments
+    pub sort_exprs: &'a [Expr], // the expressions of `order by`


If you made this Option<&'a [Expr]> that would probably be easier for UDF implementors to check if ORDER BY was specified

I think we can have an empty vec to describe no ORDER BY is given.

Another approach is always to set ordering with another function with_ordering(sort_exprs: &'a [Expr]).

alamb · 2024-04-01T20:53:13Z

datafusion/expr/src/udaf.rs

+    ///
+    /// `arg`: the type of the argument to this accumulator
+    ///
+    /// `sort_exprs`: contains a list of `Expr::SortExpr`s if the


I think these comments are now out of date

alamb · 2024-04-01T20:54:27Z

datafusion/expr/src/udaf.rs

+        &self,
+        _name: &str,
+        _value_type: DataType,
+        _ordering_fields: Vec<Field>,


What are ordering_fields used for? I think they should be documented

alamb · 2024-04-01T20:54:39Z

datafusion/expr/src/udaf.rs


    /// Return the type used to serialize the  [`Accumulator`]'s intermediate state.
    /// See [`Accumulator::state()`] for more details
    fn state_type(&self, return_type: &DataType) -> Result<Vec<DataType>>;

+    /// Return the fields of the intermediate state. It is mutually exclusive with [`Self::state_type`].


If we are going to change the AggregateUDFImpl anyways, maybe we should just remove state_type and always require state_fields?

Signed-off-by: jayzhan211 <[email protected]>

alamb · 2024-04-02T18:34:01Z

datafusion/optimizer/src/analyzer/type_coercion.rs

        let accumulator: AccumulatorFactoryFunction =
            Arc::new(|_| Ok(Box::<AvgAccumulator>::default()));
        let my_avg = AggregateUDF::from(SimpleAggregateUDF::new_with_signature(
            "MY_AVG",
            Signature::uniform(1, vec![DataType::Float64], Volatility::Immutable),
            return_type,
            accumulator,
-            state_type,
+            vec![
+                Field::new("count", DataType::UInt64, true),


I actually think this is a nice change that the fields are now named rather than anonymous

I think @jacksonrnewhouse mentioned this might be helpful for his project as well

alamb

Thanks @jayzhan211 -- I think this looks good to me

What are you thinking of next? To try and pull first_value into its own crate (to prove it is separable)?

jayzhan211 · 2024-04-03T00:11:43Z

Thanks @jayzhan211 -- I think this looks good to me

What are you thinking of next? To try and pull first_value into its own crate (to prove it is separable)?

Yes, I want to move it to its own crate and be able to register like scalar functions. And, move PhysicalSorting to physical-expr-core.

alamb · 2024-04-03T10:16:11Z

datafusion-examples/examples/advanced_udaf.rs

@@ -85,13 +87,21 @@ impl AggregateUDFImpl for GeoMeanUdaf {
    /// is supported, DataFusion will use this row oriented
    /// accumulator when the aggregate function is used as a window function
    /// or when there are only aggregates (no GROUP BY columns) in the plan.
-    fn accumulator(&self, _arg: &DataType) -> Result<Box<dyn Accumulator>> {
+    fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {


I was thinking about the impact on this API for UDAF writers last night.

Specifically, about the many existing UDAFs that exist / will exist at the time this change gets released and on the first time people encounter / try to use this API. i think the args with datatypes is much easier to use (and has less mental gymnastics to use). Thus I am going to propose an easier / beginner API for this that will require fewer changes to existing UDAFs and will be easier to use for first timers

This is what i came up with: #9920

jayzhan211 added 30 commits March 7, 2024 21:37

first draft

b94f70f

Signed-off-by: jayzhan211 <[email protected]>

clippy fix

c743d13

Signed-off-by: jayzhan211 <[email protected]>

cleanup

3a7e965

Signed-off-by: jayzhan211 <[email protected]>

use one vector for ordering req

4917f56

Signed-off-by: jayzhan211 <[email protected]>

add sort exprs to accumulator

c9e8641

Signed-off-by: jayzhan211 <[email protected]>

clippy

3a5f0d1

Signed-off-by: jayzhan211 <[email protected]>

cleanup

a3ea00a

Signed-off-by: jayzhan211 <[email protected]>

fix doc test

f349f21

Signed-off-by: jayzhan211 <[email protected]>

change to ref

6fcdaac

Signed-off-by: jayzhan211 <[email protected]>

fix typo

c3512a6

Signed-off-by: jayzhan211 <[email protected]>

fix doc

092d46e

Signed-off-by: jayzhan211 <[email protected]>

fmt

8592e6b

Signed-off-by: jayzhan211 <[email protected]>

move schema and logical ordering exprs

0f8fc24

Signed-off-by: jayzhan211 <[email protected]>

remove redudant info

3185f9f

Signed-off-by: jayzhan211 <[email protected]>

rename

3ecc772

Signed-off-by: jayzhan211 <[email protected]>

cleanup

faadc63

Signed-off-by: jayzhan211 <[email protected]>

add ignore nulls

7e33910

Signed-off-by: jayzhan211 <[email protected]>

Merge remote-tracking branch 'upstream/main' into udf-order-2

cfffcbf

fix conflict

6aaa15c

Signed-off-by: jayzhan211 <[email protected]>

backup

b74b7d2

Signed-off-by: jayzhan211 <[email protected]>

complete return_type

263e6cb

Signed-off-by: jayzhan211 <[email protected]>

complete replace

0a77e4f

Signed-off-by: jayzhan211 <[email protected]>

split to first value udf

7b26377

Signed-off-by: jayzhan211 <[email protected]>

replace accumulator

4bfd91d

Signed-off-by: jayzhan211 <[email protected]>

fmt

7f54141

Signed-off-by: jayzhan211 <[email protected]>

cleanup

6339535

Signed-off-by: jayzhan211 <[email protected]>

small fix

33ae6ee

Signed-off-by: jayzhan211 <[email protected]>

remove ordering types

b4eb865

Signed-off-by: jayzhan211 <[email protected]>

make state fields more flexible

d8ab6c5

Signed-off-by: jayzhan211 <[email protected]>

cleanup

a3bff42

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 added 3 commits March 31, 2024 22:54

add comments

5b9625f

Signed-off-by: jayzhan211 <[email protected]>

fmt

d5c3f6f

Signed-off-by: jayzhan211 <[email protected]>

fix docs

dc9549a

Signed-off-by: jayzhan211 <[email protected]>

mustafasrepo reviewed Apr 1, 2024

View reviewed changes

datafusion/core/src/execution/context/mod.rs Outdated Show resolved Hide resolved

mustafasrepo reviewed Apr 1, 2024

View reviewed changes

datafusion/core/src/execution/context/mod.rs Outdated Show resolved Hide resolved

mustafasrepo reviewed Apr 1, 2024

View reviewed changes

datafusion/core/src/physical_planner.rs Outdated Show resolved Hide resolved

jayzhan211 added 2 commits April 1, 2024 19:28

Merge remote-tracking branch 'upstream/main' into first-value-udf

7ce3d41

use exprs utils

49b4a76

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 marked this pull request as ready for review April 1, 2024 12:28

jayzhan211 requested review from alamb and mustafasrepo April 1, 2024 12:28

alamb mentioned this pull request Apr 1, 2024

DataFusion weekly project plan (Andrew Lamb) - April 1, 2024 #9899

Closed

7 tasks

alamb added the api change Changes the API exposed to users of the crate label Apr 1, 2024

alamb changed the title ~~Make FirstValue an UDAF~~ Make FirstValue an UDAF, Change AggregateUDFImpl::accumulator signature Apr 1, 2024

alamb changed the title ~~Make FirstValue an UDAF, Change AggregateUDFImpl::accumulator signature~~ Make FirstValue an UDAF, Change AggregateUDFImpl::accumulator signature, support ORDER BY for UDAFs Apr 1, 2024

alamb reviewed Apr 1, 2024

View reviewed changes

jayzhan211 added 2 commits April 2, 2024 08:55

rm state type

d70cce5

Signed-off-by: jayzhan211 <[email protected]>

add comment

29c4018

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 requested a review from alamb April 2, 2024 14:31

alamb reviewed Apr 2, 2024

View reviewed changes

alamb approved these changes Apr 2, 2024

View reviewed changes

jayzhan211 merged commit dfd4442 into apache:main Apr 3, 2024
24 checks passed

alamb reviewed Apr 3, 2024

View reviewed changes

This was referenced Apr 3, 2024

Improve AggregateUDFImpl::state_fields documentation #9919

Merged

Minor: Improve documentation for AggregateUDFImpl::accumulator and AccumulatorArgs #9920

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make FirstValue an UDAF, Change `AggregateUDFImpl::accumulator` signature, support ORDER BY for UDAFs #9874

Make FirstValue an UDAF, Change `AggregateUDFImpl::accumulator` signature, support ORDER BY for UDAFs #9874

jayzhan211 commented Mar 30, 2024 •

edited

Loading

alamb left a comment

alamb Apr 1, 2024

alamb Apr 1, 2024

alamb Apr 1, 2024

alamb Apr 1, 2024

alamb Apr 1, 2024

alamb Apr 1, 2024

jayzhan211 Apr 2, 2024

jayzhan211 Apr 2, 2024 •

edited

Loading

alamb Apr 1, 2024

alamb Apr 1, 2024

alamb Apr 1, 2024

alamb Apr 2, 2024

alamb left a comment

jayzhan211 commented Apr 3, 2024 •

edited

Loading

alamb Apr 3, 2024

alamb Apr 3, 2024

Make FirstValue an UDAF, Change AggregateUDFImpl::accumulator signature, support ORDER BY for UDAFs #9874

Make FirstValue an UDAF, Change AggregateUDFImpl::accumulator signature, support ORDER BY for UDAFs #9874

Conversation

jayzhan211 commented Mar 30, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jayzhan211 Apr 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

jayzhan211 commented Apr 3, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Make FirstValue an UDAF, Change `AggregateUDFImpl::accumulator` signature, support ORDER BY for UDAFs #9874

Make FirstValue an UDAF, Change `AggregateUDFImpl::accumulator` signature, support ORDER BY for UDAFs #9874

jayzhan211 commented Mar 30, 2024 •

edited

Loading

jayzhan211 Apr 2, 2024 •

edited

Loading

jayzhan211 commented Apr 3, 2024 •

edited

Loading